#!/usr/bin/env python

from math import pi, e, sin, cos, sqrt
from numpy import linspace,zeros,savetxt,meshgrid
from scipy.linalg import solve
from collocation_2D import lhs, rhs, bc_left, bc_right, bc_top, bc_bottom, Dirichlet
from bernstein import interpolant, interpolant2D, basis_fun_eval, basis_fun_der
from exact_2D import exact_solution
from matplotlib import pyplot as plt

import pylab as p
import mpl_toolkits.mplot3d.axes3d as p3
from matplotlib import cm

# Bernstein polynomial order - for now the same order in both directions
n = 20 
m = 20
nvar = (n+1)*(m+1) # all nodes counted for the case with non-homogenous BCs

# Solution interval [x1,x2] x [y1,y2]
x1 = 0.   
x2 = 1. 

y1 = 0.
y2 = 1. 

nd = linspace(x1,x2,n+1) # nodes
ndc = linspace(x1,x2,41)
# ndy = linspace(y1,y2,m+1) # nodes in y direction

# Bernstein polynomial coefficients - the principal unknown
beta = zeros((n+1,m+1)) # This way (by setting zeros) homogenous BC will be enforced as well,
                        # In the sequel we treat only inner points if the problem has homogenous BC's.

#print beta

f = zeros(nvar) # RHS vector 
for i in range(0,n+1):
    x = nd[i]
    for j in range(0,m+1):
        y = nd[j]
        node = (m+1)*(i-1)+j  # Changed from (m-1) to (m+1) 
#       Non-homogenous BCs...
        if (i==0):
#       Left side:: Run trough all betas
	    f[node] = bc_left(x,y)    
        elif (j==0):
#       Bottom:: Run trough all betas
	    f[node] = bc_bottom(x,y) 
        elif (i==n):
#       Right: Run trough all betas
	    f[node] = bc_right(x,y) 
        elif (j==n):
#       Top: Run trough all betas
	    f[node] = bc_top(x,y) 
        else:  
	    f[node] = rhs(x,y)  

#print f

K = zeros( (nvar,nvar) ) # LHS matrix 
for i in range(0,n+1):
    x = nd[i]
    for j in range(0,m+1):
        y = nd[j]
        node = (m+1)*(i-1)+j # node defines specific location on a grid
#       Non-homogenous BCs...
        if (i==0):
#       Left side:: Run trough all betas
             for k in range(0,n+1):
                for l in range(0,m+1):
                    jfun = (m+1)*(k-1)+l
	            K[node,jfun] = Dirichlet(l,k,n,m,x1,x2,y1,y2,x,y)    
        elif (j==0):
#       Bottom:: Run trough all betas
            for k in range(0,n+1):
                for l in range(0,m+1):
                    jfun = (m+1)*(k-1)+l
	            K[node,jfun] = Dirichlet(l,k,n,m,x1,x2,y1,y2,x,y) 
        elif (i==n):
#       Right: Run trough all betas
            for k in range(0,n+1):
                for l in range(0,m+1):
                    jfun = (m+1)*(k-1)+l
	            K[node,jfun] = Dirichlet(l,k,n,m,x1,x2,y1,y2,x,y) 
        elif (j==n):
#       Top: Run trough all betas
            for k in range(0,n+1):
                for l in range(0,m+1):
                    jfun = (m+1)*(k-1)+l
	            K[node,jfun] = Dirichlet(l,k,n,m,x1,x2,y1,y2,x,y) 
        else:
#       Interior: Run trough all betas
            for k in range(0,n+1):
                for l in range(0,m+1):
                    jfun = (m+1)*(k-1)+l
	            K[node,jfun] = lhs(l,k,n,m,x1,x2,y1,y2,x,y)  

#print K

# Check sparsity pattern of matrix K:
#plt.title('Sparsity pattern of discretization matrix')
#plt.spy(K, precision=1e-10, marker='s', markersize=5)
#plt.show()

c = zeros(nvar) # 1d array holding all beta coefficients
c[0:nvar] = solve(K,f)  # Solve linear system

#print 'Solution: c'
#print c

# NOTE: c is auxilliary 1d array - we now transfer those values to 2d array beta
for i in range(0,n+1):
    for j in range(0,m+1):
        node = (m+1)*(i-1)+j # note the way it traverses the gird - first up, then go to next column on the right and up again.
        beta[j][i] = c[node]

#print 'Solution matrix beta:'
savetxt('solution-beta.txt', beta)
#print beta

u = zeros((n+1,m+1))
#u = zeros((ndc,ndc))
for i in range(n+1):
	x = nd[i]
        for j in range(m+1):
            y = nd[j]
	    u[i][j] = interpolant2D(n,m,x1,x2,y1,y2,beta,x,y)  # Solution - a Bernstein interpolant constructed using beta's
                                                               # at whatever points - we use collocation points here.
# Save solution matrix to a ASCII file:
savetxt('solution.txt', u)
#print u

uex = zeros((n+1,m+1))
for i in range(n+1):
	x = nd[i]
        for j in range(m+1):
            y = nd[j]
	    uex[i][j] = exact_solution(x,y)

# Print absolute error
#print uex
savetxt('solution-exact.txt', uex)
savetxt('solution-err.txt', abs(u-uex))

# L2 relative error norm
suma1 = 0.
suma2 = 0.
for i in range(n+1):
        for j in range(m+1):
	    suma1 += (u[i][j]-uex[i][j])**2
            suma2 += uex[i][j]**2
l2errnrm = sqrt(suma1/suma2)
print l2errnrm

# wireframe:
fig=p.figure()
ax = p3.Axes3D(fig)
X, Y = meshgrid(nd, nd)
ax.plot_wireframe(X,Y,u)
ax.plot_surface(X,Y,u,rstride=1, cstride=1, alpha=1.0, cmap=cm.cool,
        linewidth=0, antialiased=False)
#ax.plot_surface(X,Y,u,rstride=1, cstride=1,alpha=0.5)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
p.show()



